package query

import (
	"fmt"
	"regexp"
	"strconv"
	"strings"
	"sync"
)

// AscendQuery Ascend NPU查询实现
type AscendQuery struct{}

func init() {
	RegisterQueryer("ascend", func() GPUQuery { return &AscendQuery{} })
}

func (q *AscendQuery) Query() ([]GPUInfo, error) {
	// 获取设备列表
	devices, err := q.getDevices()
	if err != nil {
		return nil, fmt.Errorf("获取设备列表失败: %v", err)
	}

	// 并行获取所有设备信息
	var wg sync.WaitGroup
	gpuInfos := make([]GPUInfo, len(devices))
	errors := make([]error, len(devices))

	for i, device := range devices {
		wg.Add(1)
		go func(index int, device int) {
			defer wg.Done()
			utilization, memoryUsed, memoryTotal, err := q.getDeviceInfo(device)
			if err != nil {
				errors[index] = err
				return
			}
			gpuInfos[index] = GPUInfo{
				Index:        device,
				Utilization: utilization,
				MemoryUsed:   memoryUsed,
				MemoryTotal:  memoryTotal,
			}
		}(i, device)
	}
	wg.Wait()

	// 检查是否有错误
	for _, err := range errors {
		if err != nil {
			return nil, fmt.Errorf("获取设备信息失败: %v", err)
		}
	}

	return gpuInfos, nil
}

// getDevices 获取设备列表
func (q *AscendQuery) getDevices() ([]int, error) {
	output, err := runCommand("npu-smi", "info", "-l")
	if err != nil {
		return nil, err
	}

	// 使用正则表达式快速提取设备ID
	deviceRegex := regexp.MustCompile(`NPU ID\s+:\s+(\d+)`)
	lines := strings.Split(output, "\n")
	var devices []int

	for _, line := range lines {
		if matches := deviceRegex.FindStringSubmatch(line); matches != nil {
			if deviceID, err := strconv.Atoi(matches[1]); err == nil {
				devices = append(devices, deviceID)
			}
		}
	}

	if len(devices) == 0 {
		return nil, fmt.Errorf("未找到有效的NPU设备")
	}

	return devices, nil
}

// getDeviceInfo 获取设备使用率和内存信息
func (q *AscendQuery) getDeviceInfo(device int) (int, int, int, error) {
	// 获取使用率信息
	usageOutput, err := runCommand("npu-smi", "info", "-t", "usages", "-i", strconv.Itoa(device))
	if err != nil {
		return 0, 0, 0, err
	}

	// 获取内存信息
	memoryOutput, err := runCommand("npu-smi", "info", "-t", "memory", "-i", strconv.Itoa(device))
	if err != nil {
		return 0, 0, 0, err
	}

	// 解析使用率信息
	utilization, err := q.parseUtilization(usageOutput)
	if err != nil {
		return 0, 0, 0, err
	}

	// 解析内存信息
	memoryUsed, memoryTotal, err := q.parseMemoryInfo(memoryOutput, usageOutput)
	if err != nil {
		return 0, 0, 0, err
	}

	return utilization, memoryUsed, memoryTotal, nil
}

// parseUtilization 解析使用率信息
func (q *AscendQuery) parseUtilization(output string) (int, error) {
	// 预编译正则表达式
	utilRegex := regexp.MustCompile(`Aicore Usage Rate\(%\)\s+:\s+(\d+)`)
	matches := utilRegex.FindStringSubmatch(output)
	if matches == nil {
		return 0, fmt.Errorf("未找到Aicore使用率信息")
	}

	if utilization, err := strconv.Atoi(matches[1]); err == nil {
		return utilization, nil
	}

	return 0, fmt.Errorf("解析使用率失败")
}

// parseMemoryInfo 解析内存信息
func (q *AscendQuery) parseMemoryInfo(memoryOutput, usageOutput string) (int, int, error) {
	// 预编译正则表达式
	memoryTotalRegex := regexp.MustCompile(`HBM Capacity\(MB\)\s+:\s+(\d+)`)
	memoryUsageRegex := regexp.MustCompile(`HBM Usage Rate\(%\)\s+:\s+(\d+)`)

	// 查找内存总量
	totalMatches := memoryTotalRegex.FindStringSubmatch(memoryOutput)
	if totalMatches == nil {
		return 0, 0, fmt.Errorf("未找到HBM总容量信息")
	}
	total, err := strconv.Atoi(totalMatches[1])
	if err != nil {
		return 0, 0, fmt.Errorf("解析内存总量失败: %v", err)
	}

	// 查找内存使用率
	usageMatches := memoryUsageRegex.FindStringSubmatch(usageOutput)
	if usageMatches == nil {
		return 0, 0, fmt.Errorf("未找到HBM使用率信息")
	}
	usageRate, err := strconv.Atoi(usageMatches[1])
	if err != nil {
		return 0, 0, fmt.Errorf("解析内存使用率失败: %v", err)
	}

	// 计算内存使用量
	used := (total * usageRate) / 100

	return used, total, nil
}
